package org.raymond.iworks.core.security;

import cn.hutool.core.codec.Base64;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.symmetric.AES;
import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;
import lombok.extern.slf4j.Slf4j;
import org.raymond.iworks.apis.common.Result;
import org.springframework.http.HttpHeaders;
import org.springframework.util.FastByteArrayOutputStream;

import java.io.*;

@Slf4j
public class CachedEncryptResponseWrapper extends HttpServletResponseWrapper {
    private AES aes;
    private final FastByteArrayOutputStream content = new FastByteArrayOutputStream(1024);
    private ServletOutputStream outputStream;
    private PrintWriter writer;
    private Integer contentLength;

    public CachedEncryptResponseWrapper(HttpServletResponse response, AES aes) {
        super(response);
        response.setCharacterEncoding("UTF-8");
        this.aes = aes;
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        if (this.outputStream == null) {
            this.outputStream = new ResponseServletOutputStream(getResponse().getOutputStream());
        }
        return this.outputStream;
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        if (this.writer == null) {
            this.writer = new ResponsePrintWriter("UTF-8");
        }
        return this.writer;
    }


    @Override
    public void flushBuffer() throws IOException {
        // do not flush the underlying response as the content as not been copied to it yet
    }

    @Override
    public void setContentLength(int len) {
        if (len > this.content.size()) {
            this.content.resize(len);
        }
        this.contentLength = len;
    }

    // Overrides Servlet 3.1 setContentLengthLong(long) at runtime
    @Override
    public void setContentLengthLong(long len) {
        if (len > Integer.MAX_VALUE) {
            throw new IllegalArgumentException("Content-Length exceeds ContentCachingResponseWrapper's maximum (" +
                    Integer.MAX_VALUE + "): " + len);
        }
        int lenInt = (int) len;
        if (lenInt > this.content.size()) {
            this.content.resize(lenInt);
        }
        this.contentLength = lenInt;
    }

    @Override
    public void setBufferSize(int size) {
        if (size > this.content.size()) {
            this.content.resize(size);
        }
    }

    @Override
    public void resetBuffer() {
        this.content.reset();
    }

    @Override
    public void reset() {
        super.reset();
        this.content.reset();
    }

    /**
     * Return the cached response content as a byte array.
     */
    public byte[] getContentAsByteArray() {
        return this.content.toByteArray();
    }

    /**
     * Return an {@link InputStream} to the cached content.
     * @since 4.2
     */
    public InputStream getContentInputStream() {
        return this.content.getInputStream();
    }

    /**
     * Return the current size of the cached content.
     * @since 4.2
     */
    public int getContentSize() {
        return this.content.size();
    }

    public void copyBodyToResponse() throws IOException {
        copyBodyToResponse(true);
    }

    protected void copyBodyToResponse(boolean complete) throws IOException {
        if (this.content.size() > 0) {
            HttpServletResponse rawResponse = (HttpServletResponse) getResponse();
            if ((complete || this.contentLength != null) && !rawResponse.isCommitted()) {
                if (rawResponse.getHeader(HttpHeaders.TRANSFER_ENCODING) == null) {
                    rawResponse.setContentLength(complete ? this.content.size() : this.contentLength);
                }
                this.contentLength = null;
            }
            this.content.writeTo(rawResponse.getOutputStream());
            this.content.reset();
            if (complete) {
                super.flushBuffer();
            }
        }
    }

    public void updateContent() throws IOException {
        byte[] oriBytes = this.content.toByteArray();
        String oriResponse = StrUtil.str(oriBytes, CharsetUtil.CHARSET_UTF_8);
        Result result = null; //JSONObject.parseObject(oriResponse, Result.class);
        String resultStr = result.getData()==null?"":result.getData().toString();
        String encBase64 = "";
        if(aes!=null){
            encBase64 = aes.encryptBase64(resultStr);
        }else{
            encBase64 = Base64.encode(resultStr);
        }
        log.info("str:{}, b64:{}", resultStr, encBase64);
        result.setData(encBase64);
        byte[] encryptResponse = null; //JSONObject.toJSONString(result).getBytes(StandardCharsets.UTF_8);
        reset();
        this.content.write(encryptResponse);
        this.content.writeTo(getResponse().getOutputStream());
        this.content.reset();
        super.flushBuffer();
    }

    private class ResponseServletOutputStream extends ServletOutputStream {
        private final ServletOutputStream os;
        public ResponseServletOutputStream(ServletOutputStream os) {
            this.os = os;
        }
        @Override
        public void write(int b) throws IOException {
            content.write(b);
        }
        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            content.write(b, off, len);
        }
        @Override
        public boolean isReady() {
            return this.os.isReady();
        }
        @Override
        public void setWriteListener(WriteListener writeListener) {
            this.os.setWriteListener(writeListener);
        }
    }


    private class ResponsePrintWriter extends PrintWriter {
        public ResponsePrintWriter(String characterEncoding) throws UnsupportedEncodingException {
            super(new OutputStreamWriter(content, characterEncoding));
        }
        @Override
        public void write(char[] buf, int off, int len) {
            super.write(buf, off, len);
            super.flush();
        }
        @Override
        public void write(String s, int off, int len) {
            super.write(s, off, len);
            super.flush();
        }
        @Override
        public void write(int c) {
            super.write(c);
            super.flush();
        }
    }
}
